#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#pragma OPENCL EXTENSION cl_arm_printf : enable
__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;
constant float Bt[36] = {
    1.0000000000f, 0.0000000000f, -2.5000004768f, -0.0000001192f, 1.0000001192f, 0.0000000000f, 
    0.0000000000f, 0.9428091049f, 1.3333333731f, -0.4714044929f, -0.6666667461f, 0.0000000000f, 
    0.0000000000f, -0.9428089857f, 1.3333334923f, 0.4714045525f, -0.6666667461f, 0.0000000000f, 
    0.0000000000f, -0.1178511307f, -0.0833333358f, 0.2357022613f, 0.1666666865f, 0.0000000000f, 
    0.0000000000f, 0.1178511307f, -0.0833333507f, -0.2357022911f, 0.1666666865f, 0.0000000000f, 
    0.0000000000f, 0.9999998808f, -0.0000000596f, -2.5000000000f, 0.0000000000f, 1.0000000000f, 
};

__kernel void winograd_4x4_to_36_image2d(__write_only image2d_t output_tex2d,
                                         __read_only image2d_t input_tex2d,
                                         __read_only image2d_t bt_tex2d,
                                         int4 input_sliced_shape,        // 1, 480, 480, 32;      1, 16, 256, 24
                                         int4 output_sliced_shape) {     // 1, 36, 14400, 32,     1, 36, 256, 24 
    int tile_id = get_global_id(0);
    int tile_row = get_global_id(1);
    int slice = get_global_id(2);

    int tile_total = output_sliced_shape.z, tile_length = output_sliced_shape.y, sliced_channel = output_sliced_shape.w;    // 14400, 36, 32
    int height = input_sliced_shape.y, width = input_sliced_shape.z;                                                        // 480, 480
    int padding = 1;

    if (tile_id >= tile_total || tile_row >= 6 || slice >= sliced_channel)   // 溢出检测
        return ;    
    
    int tiles_in_width = width / 4;                         // 120
    int tile_x_idx = (tile_id % tiles_in_width) * 4;        // 获取一个tile，此处计算这个tile左上角的坐标
    int tile_y_idx = (tile_id / tiles_in_width) * 4;

    int x_idx[6];
    bool x_idx_valid[6];
    for (int i=0; i<6; i++) {
        x_idx[i] = tile_x_idx - padding + i;      // 获取这个tile在slice内部的局部的x方向坐标
        x_idx_valid[i] = (x_idx[i] >= 0 && x_idx[i] < width);
        x_idx[i] = clamp(x_idx[i], 0, width-1);
    }

    float4 bt_row_tile_row_0 = convert_float4(read_imagef(bt_tex2d, smp_none, (int2)(tile_row*2 + 0, 0)));  // 读取bt_matrix的第tile_row行，即两个float4
    float4 bt_row_tile_row_1 = convert_float4(read_imagef(bt_tex2d, smp_none, (int2)(tile_row*2 + 1, 0)));
    float bt_row[6] = {bt_row_tile_row_0.x, bt_row_tile_row_0.y, bt_row_tile_row_0.z,                        // prepare bt
                       bt_row_tile_row_0.w, bt_row_tile_row_1.x, bt_row_tile_row_1.y};
    float4 mul_intermidate[6] = {0};
    int yoffset = slice * height;
    for (int i=0; i<6; i++) {
        int y_idx_in_slice = tile_y_idx - padding + i;                      // 获取这个tile在slice内部的局部的y方向坐标
        bool y_idx_valid = (y_idx_in_slice >= 0 && y_idx_in_slice < height);
        int y_idx = yoffset + clamp(y_idx_in_slice, 0, height-1);
        // int offset = select(0, y_idx_in_slice * width, y_idx_valid);        // 找到这个y_idx在slice内部的全局坐标
        float bt_val = bt_row[i] * (float)(y_idx_valid);                      // 获取一个bt
        for (int j=0; j<6; j++) {
            // float4 input_val = input_tex2d[x_idx[j] + offset] * (float)(x_idx_valid[j]);
            float4 input_val = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(x_idx[j], y_idx))) * (float)(x_idx_valid[j]);
            mul_intermidate[j] += bt_val * input_val;
        }
    }

    tile_row = tile_row * 6;
    for (int i=0; i<6; i++) {
        float4 accVal = 0;
        for (int j=0; j<6; j++) {
            accVal += Bt[6*i+j]*mul_intermidate[j];
        }
        write_imagef(output_tex2d, (int2)(tile_id, slice*tile_length+tile_row), accVal);
        // output_tex2d[(slice*tile_length+tile_row)*tile_total+tile_id] = accVal; // (32*(36*(14400*4)))
        tile_row = tile_row + 1;
    }
}